{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "W1_Tutorial1",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernel": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.8"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/CIS-522/course-content/blob/main/tutorials/W01_AlphaZero/W1_Tutorial1.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ML2DwqwkVwfo"
      },
      "source": [
        "# CIS-522 Week 1 Part 1\n",
        "# Introduction and AlphaZero\n",
        "\n",
        "__Instructor:__ Konrad Kording\n",
        "\n",
        "__Content creators:__ Ari Benjamin, Richard Lange, Byron Galbraith\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-fhwfTfcPKzo"
      },
      "source": [
        "---\n",
        "# Course introduction\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5eXcDCI3YUdl"
      },
      "source": [
        "## Know your tool: Google Colab."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4HWKgQ5rcvnd"
      },
      "source": [
        "This environment is a Colab. It's a Jupyter notebook hosted on the Google Cloud. We will do everything here: embed videos, embed questionnaires, and run our code.\n",
        "\n",
        "You'll soon be experts at this. One thing to remember is **always run all cells!**. We need you to press `Shift+Enter` or press the play button to show the videos and save your responses."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_MpOtSC-bsnH",
        "cellView": "form"
      },
      "source": [
        "#@markdown What is your Pennkey and pod? (text, not numbers, e.g. bfranklin)\n",
        "my_pennkey = 'value' #@param {type:\"string\"}\n",
        "my_pod = 'sublime-newt' #@param ['Select', 'euclidean-wombat', 'sublime-newt', 'buoyant-unicorn', 'lackadaisical-manatee','indelible-stingray','superfluous-lyrebird','discreet-reindeer','quizzical-goldfish','astute-jellyfish','ubiquitous-cheetah','nonchalant-crocodile','fashionable-lemur','spiffy-eagle','electric-emu','quotidian-lion']\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hpZ6TBJq554D"
      },
      "source": [
        "#### Ensure you're running a GPU notebook.\n",
        "\n",
        "From \"Runtime\" in the drop-down menu above, click \"Change runtime type\". Ensure that \"Hardware Accelerator\" says \"GPU\".\n",
        "\n",
        "#### Ensure you can save!\n",
        "\n",
        "From \"File\", click \"Save a copy in Drive\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zhERpJAhOOcZ"
      },
      "source": [
        "#### Save your answers\n",
        "\n",
        "After you type into a form, like you just did, make sure you ran the cell to save your answer. Your answer will be saved in a variable:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WU4CeGgiOS3V"
      },
      "source": [
        "print(my_pennkey)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KUTMCuCARMKc"
      },
      "source": [
        "## Who we are"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jIIUpuZJSWx2",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Lyle and Konrad\n",
        "from IPython.display import YouTubeVideo\n",
        "video = YouTubeVideo(id=\"4DuhPv8EsWk\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "import time\n",
        "try: t0;\n",
        "except NameError: t0=time.time()\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "v8XJu__JR_qv",
        "cellView": "form"
      },
      "source": [
        "#@markdown (RUN ME) Show student locations. Drag and zoom!\n",
        "lat = [39.95, 40.35, 39.95, 32.7, 39.95, 39.95, 39.95, 39.95, 39.95, 39.95, 39.95, 41.03, 39.95, 40.41, 39.95, 39.95, 39.95, 39.95, 12.98, 39.95, 40.32, 39.95, 39.95, 39.95, 39.95, 39.95, 19.08, 39.95, 39.95, 39.95, 39.95, 39.95, 39.95, 39.95, 39.95, 39.95, 49.28, 37.34, 39.95, 39.95, 32.72, 45.63, 39.95, 39.95, 39.95, 39.95, 39.95, 39.95, 39.95, 39.95, 39.95, 32.06, 39.95, 39.95, 44.02, 39.95, 40.15, 28.0, 39.95, 39.95, 38.92, 47.48, 39.75, 22.56, 39.95, 22.28, 39.95, 39.95, 39.95, 39.95, 37.34, 39.95, 39.95, 49.26, 48.14, 39.95, 37.0, 37.57, 42.38, 39.95, 39.95, 39.95, 28.71, 39.95, 39.8, 39.95, 39.95, 39.95, 39.95, 37.32, 39.95, 39.95, 39.95, 39.95, 35.0, 39.95, 38.97, 43.59, 39.95, 39.95, 32.78, 30.33, 30.6, 39.91, 39.95, 40.54, 39.95, 39.95, 49.26, 39.95, 38.92, 39.95, 39.96, 39.95, 38.93, 39.95, 39.95, 39.95, 42.36, 39.95, 39.95, 39.95, 39.95, 41.12, 39.95, 34.34, 39.95, 39.95, 39.95, 39.95, 39.95, 39.95, 39.1, 39.95, 39.95, 39.95, 43.07, 30.25, 39.95, 39.95]\n",
        "lon = [-75.16, -74.66, -75.16, -97.11, -75.16, -75.16, -75.16, -75.16, -75.16, -75.16, -75.16, -74.1, -75.16, -74.57, -75.16, -75.16, -75.16, -75.16, 77.59, -75.16, -74.62, -75.16, -75.16, -75.16, -75.16, -75.16, 72.88, -75.16, -75.16, -75.16, -75.16, -75.16, -75.16, -75.16, -75.16, -75.16, -122.79, -121.89, -75.16, -75.16, -117.16, -122.67, -75.16, -75.16, -75.16, -75.16, -75.16, -75.16, -75.16, -75.16, -75.16, 118.79, -75.16, -75.16, -92.46, -75.16, -74.71, 120.7, -75.16, -75.16, 121.63, 19.15, -75.55, 114.05, -75.16, 114.16, -75.16, -75.16, -75.16, -75.16, -121.89, -75.16, -75.16, -123.11, 11.58, -75.16, 112.0, 126.98, -71.11, -75.16, -75.16, -75.16, 77.44, -75.16, -75.46, -75.16, -75.16, -75.16, -75.16, -122.03, -75.16, -75.16, -75.16, -75.16, 105.0, -75.16, -76.86, -79.64, -75.16, -75.16, -96.8, -81.66, 114.3, 116.39, -75.16, -74.39, -75.16, -75.16, -123.11, -75.16, 121.63, -75.16, -75.61, -75.16, -74.92, -75.16, -75.16, -75.16, -71.06, -75.16, -75.16, -75.16, -75.16, -73.41, -75.16, 108.94, -75.16, -75.16, -75.16, -75.16, -75.16, -75.16, -76.85, -75.16, -75.16, -75.16, -89.38, 120.21, -75.16, -75.16]\n",
        "\n",
        "import plotly as py\n",
        "import plotly.graph_objects as go\n",
        "import numpy as np\n",
        "\n",
        "fig = go.Figure()\n",
        "\n",
        "fig.add_trace(\n",
        "    go.Scattergeo(\n",
        "        locationmode = 'country names',\n",
        "        lon = np.array(lon) + np.random.randn(len(lon))*1e-2,\n",
        "        lat = np.array(lat) + np.random.randn(len(lon))*1e-2,\n",
        "        marker = {\"color\" :'red'}\n",
        "    )\n",
        ")\n",
        "\n",
        "fig.update_layout(\n",
        "    title_text = '',\n",
        "    showlegend = False,\n",
        "    geo = go.layout.Geo(\n",
        "        scope = 'world',\n",
        "        projection_type = 'natural earth',lakecolor = 'rgb(199, 226, 255)',\n",
        "        showland = True, showocean = True,showlakes = True,\n",
        "        landcolor = 'rgb(245, 255, 230)',\n",
        "        countrycolor = 'rgb(204, 204, 204)',\n",
        "        oceancolor='rgb(199, 226, 255)'\n",
        "    ),\n",
        ")\n",
        "\n",
        "py.offline.iplot(fig)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "z-BSM9Xpcyrv",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Welcome to the Deep Learning Academy\n",
        "video = YouTubeVideo(id=\"dFpprmYImyM\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LCExULorcT2r",
        "cellView": "form"
      },
      "source": [
        "#@markdown Which week are you most excited for?\n",
        "week = 0 #@param {type:\"number\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ijiuk5F6iKKo"
      },
      "source": [
        "## Course logistics\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Pbcqy6CCWGG4",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Course Logistics\n",
        "\n",
        "param_list = ['week']\n",
        "for param in param_list:\n",
        "    if param not in locals():\n",
        "        raise NameError(\"Please make sure to run the cell \"\n",
        "                        \"after submitting your answer for <{}>!\".format(param))\n",
        "\n",
        "video = YouTubeVideo(id=\"yf_PZcPumOo\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BJr6cxKyiJHK",
        "cellView": "form"
      },
      "source": [
        "#@markdown How important is learning theory vs. implementation skills for you?\n",
        "theory_or_implementation = 5.7 #@param {type:\"slider\", min:0, max:10, step:0.1}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fVqMxr6fOU3z"
      },
      "source": [
        "Let 0 be all theory and no implementation, and 10 be all implementation and no theory.\n",
        "\n",
        "Again, make sure to run the cell to save your answer."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "duA79yfviP2H"
      },
      "source": [
        "## Interacting with this course"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "64Ak82pgiSU9",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Standards for Conduct\n",
        "\n",
        "param_list = ['theory_or_implementation']\n",
        "for param in param_list:\n",
        "    if param not in locals():\n",
        "        raise NameError(\"Please make sure to run the cell \"\n",
        "                        \"after submitting your answer for <{}>!\".format(param))\n",
        "\n",
        "video = YouTubeVideo(id=\"iJ7mwRRgGSk\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sz8K4qdo9lLh"
      },
      "source": [
        "After you watch the video, please read through the [code of conduct](https://docs.google.com/document/d/1eHKIkaNbAlbx_92tLQelXnicKXEcvFzlyzzeWjEtifM/edit?usp=sharing)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lRzrCfj0MuEC"
      },
      "source": [
        "---\n",
        "# BREAK. \n",
        "Take 5 minutes and get some tea."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "14u5mUEFUA3B"
      },
      "source": [
        "---\n",
        "# Setup"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kTyWn2EzT_pd"
      },
      "source": [
        "# Imports\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "from torch import nn\n",
        "import numpy as np\n",
        "import os\n",
        "from numba import cuda\n",
        "from matplotlib import pyplot as plt\n",
        "import matplotlib.patches as patches\n",
        "from typing import List\n",
        "from copy import deepcopy\n",
        "\n",
        "from tqdm.notebook import tqdm\n",
        "import IPython\n",
        "from IPython.display import IFrame\n",
        "from google.colab import output\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Un34UwDpUkre",
        "cellView": "form"
      },
      "source": [
        "# @title Figure Settings\n",
        "%config InlineBackend.figure_format = 'retina'\n",
        "%matplotlib inline \n",
        "\n",
        "fig_w, fig_h = (8, 6)\n",
        "plt.rcParams.update({'figure.figsize': (fig_w, fig_h)})\n",
        "\n",
        "plt.style.use(\"https://raw.githubusercontent.com/NeuromatchAcademy/\"\n",
        "              \"course-content/master/nma.mplstyle\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-CKeaaebUr9L",
        "cellView": "form"
      },
      "source": [
        "# @title Helper functions\n",
        "\n",
        "def outcome2color(outcome, alpha=1):\n",
        "    \"\"\"Return RGB color for given outcome (-1, 0, +1)\n",
        "    \"\"\"\n",
        "    if outcome == +1:\n",
        "        return (40/255., 229./255., 40/255., alpha) # Win as green\n",
        "    elif outcome == -1:\n",
        "        return (229./255., 40/255., 0, alpha) # Loss as red\n",
        "    else:\n",
        "        return (0, 76/255., 220/255., alpha) # Draw as blue\n",
        "\n",
        "def win_statistics(agent1, agent2, games_per_side=100, bar_plot=True):\n",
        "    \"\"\"Play agent1 vs agent2 many times to get some statistics. At the end, show\n",
        "    statistics with a bar plot and return wins/draws/losses for agent1. Runs\n",
        "    'games_per_side' matches with agent1 as player 1 and another 'games_per_side'\n",
        "    matches with agent1 as player 2.\n",
        "\n",
        "    The returned 'win_info' table is size 2x3. Row 0 counts agent1 as player 1.\n",
        "    Row 1 counts agent1 as player 2. Col 0 is wins for agent1, col 1 is draws,\n",
        "    and col 2 is losses\n",
        "    \"\"\"\n",
        "    win_info = np.zeros((1, 3))\n",
        "    # Play agent1 as player 1\n",
        "    outcome1, _ = ai_vs_ai(agent1, agent2, n_games=games_per_side//2, \n",
        "                          n=6)\n",
        "    # Play agent1 as player 2\n",
        "    outcome2, _ = ai_vs_ai(agent2, agent1, n_games=games_per_side//2, \n",
        "                          n=6)\n",
        "    outcome = torch.cat((outcome1,outcome2),dim=0) \n",
        "    for g in range(games_per_side):\n",
        "        col_idx = 1-outcome[g,0] # Map agent1 outcome (-1,0,1) to column index (2,1,0)\n",
        "        win_info[0, col_idx] += 1\n",
        "    # (Maybe) plot\n",
        "    if bar_plot:\n",
        "        # Approximate standard error in each probability estimate (it would be\n",
        "        # more correct to do confidence intervals on the binomial distribution\n",
        "        # for each of win/draw/loss, but this is fine for a large # of games)\n",
        "        win_prob = win_info / games_per_side\n",
        "        win_se = np.sqrt(win_prob*(1-win_prob)/games_per_side)\n",
        "        plt.figure(figsize=(4,6))\n",
        "        plt.bar([0], win_prob[:,0], yerr=win_se[:,0], color=outcome2color(1, .8), label=\"wins\")\n",
        "        plt.bar([0], win_prob[:,1], yerr=win_se[:,1], bottom=win_prob[:,0], color=outcome2color(0, 0.8))\n",
        "        plt.bar([0], win_prob[:,2], yerr=win_se[:,2], bottom=win_prob[:,:2].sum(axis=1), color=outcome2color(-1, 0.8))\n",
        "        plt.xticks([])\n",
        "        plt.yticks(np.linspace(0,1,11))\n",
        "        plt.ylabel(\"Fraction games won\")\n",
        "        plt.title(f\"{agent1.name} vs {agent2.name}\")\n",
        "        plt.legend([f\"{agent1.name} wins\", \"draws\", f\"{agent1.name} loses\"])\n",
        "        plt.grid()\n",
        "        plt.show()\n",
        "    return win_info\n",
        "\n",
        "\n",
        "def actions_as_tensor(actions_list):\n",
        "  \"\"\"Takes a list of actions and returns the same as a Tensor.\n",
        "  Handles the case when some games have more available actions.\n",
        "  Inputs:\n",
        "     actions_list is a list of lists of tuples. One list per game.\n",
        "\n",
        "  \"\"\"\n",
        "  num_actions_each_game = torch.as_tensor(\n",
        "                           [len(acts) for acts in actions_list])\n",
        "  max_actions = num_actions_each_game.max()\n",
        "  n_games = len(actions_list)\n",
        "  values = float('-inf')*torch.ones(n_games, max_actions)\n",
        "  # Create a tensor of [n_games, max_actions, 2] that contains\n",
        "  # all available actions for all games, padded with -1s (pass move)\n",
        "  # where games have fewer than max_actions possible actions.\n",
        "  actions_tensor = -torch.ones(n_games, max_actions, 2,\n",
        "                                dtype=torch.int32)\n",
        "  for g in range(game.n_games):\n",
        "      if num_actions_each_game[g] == 0:\n",
        "          continue\n",
        "      actions_tensor[g, :num_actions_each_game[g], :] =\\\n",
        "          torch.as_tensor(actions_list[g], dtype=torch.int32)\n",
        "  return actions_tensor\n",
        "# "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E4aG-P9lfeg7"
      },
      "source": [
        "---\n",
        "# Section 1: implementing AlphaZero"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4QFiTOwR7zk1",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Lesson 1 Intro + Othello\n",
        "video = YouTubeVideo(id=\"iWNkBeK8WZ0\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "try: t1;\n",
        "except NameError: t1=time.time()\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bcavjkibRsLJ"
      },
      "source": [
        "### Objectives for today\n",
        " - Understand Othello as we build up to solving it with AlphaZero\n",
        " - Understand how a NN can be used in gameplay\n",
        " - Build examples and objectives\n",
        " - Train our first neural network\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KSNVIY2tilRl",
        "cellView": "form"
      },
      "source": [
        "#@title Section 1.1: Create OthelloGame and InteractiveOthelloGame\n",
        "\n",
        "# Environment config so that numba knows where to find CUDA libraries\n",
        "os.environ['NUMBAPRO_LIBDEVICE'] = \"/usr/local/cuda-10.0/nvvm/libdevice\"\n",
        "os.environ['NUMBAPRO_NVVM'] = \"/usr/local/cuda-10.0/nvvm/lib64/libnvvm.so\"\n",
        "\n",
        "# 1st dimension is row, 2nd is col. Top-left corner of board is (0,0). 8 vector\n",
        "# directions out from each point are:\n",
        "RAYS = np.array([[0, 1],   # east\n",
        "                 [0, -1],  # west\n",
        "                 [1, 0],   # south\n",
        "                 [-1, 0],  # north\n",
        "                 [1, 1],   # southeast\n",
        "                 [1, -1],  # southwest\n",
        "                 [-1, 1],  # northeast\n",
        "                 [-1, -1]],# northwest\n",
        "                dtype=np.int32)\n",
        "\n",
        "# ======================================\n",
        "# === BEGIN CUDA-ACCELERATED HELPERS ===\n",
        "# ======================================\n",
        "\n",
        "def torch2cuda(arr:torch.Tensor) -> cuda.cudadrv.devicearray.DeviceNDArray:\n",
        "    return cuda.as_cuda_array(arr)\n",
        "\n",
        "def cuda2torch(arr:cuda.cudadrv.devicearray.DeviceNDArray) -> torch.Tensor:\n",
        "    return torch.as_tensor(arr, device='cuda')\n",
        "\n",
        "@cuda.jit(device=True)\n",
        "def cuda_cast_rays(board, y0, x0, ray, player):\n",
        "    \"\"\"Helper CUDA kernel that searches out from row, column coordinate (y0,x0)\n",
        "    in the direction specified by \"ray\". Returns True iff there is an unbroken\n",
        "    line of opponent stones along the ray starting at (y0,x0), ending with a\n",
        "    stone belonging to \"player\".\n",
        "    \n",
        "    The device=True decorator means that this may only be called from other CUDA\n",
        "    kernels (see below).\n",
        "    \"\"\"\n",
        "    n = board.shape[1]\n",
        "    opponent = 3 - player\n",
        "    y, x = y0+ray[0], x0+ray[1]\n",
        "\n",
        "    # The ray must begin with an opponent stone 1 space away\n",
        "    if x < 0 or y < 0 or x >= n or y >= n or board[y, x] != opponent:\n",
        "        return False\n",
        "    \n",
        "    # Keep looking until edge of the board or we find an open space or we find\n",
        "    # a player-owned stone\n",
        "    y, x = y+ray[0], x+ray[1]\n",
        "    while x >= 0 and y >= 0 and x < n and y < n:\n",
        "        # Ray terminates on an open space. Not a valid move.\n",
        "        if board[y, x] == 0:\n",
        "            return False\n",
        "        \n",
        "        # Ray terminates on player. This is valid!\n",
        "        if board[y, x] == player:    \n",
        "            return True\n",
        "        \n",
        "        y, x = y+ray[0], x+ray[1]\n",
        "        \n",
        "    # Reached end of the board. Not a valid move.\n",
        "    return False\n",
        "\n",
        "@cuda.jit\n",
        "def cuda_get_valid_actions(boards, player, valid_actions):\n",
        "    \"\"\"Helper CUDA kernel that calls cuda_cast_rays to find all legal moves\n",
        "    in parallel across many boards and positions. Results are written into the\n",
        "    array \"valid_actions\" which must be the same shape as \"boards\".\n",
        "\n",
        "    Decorating with @cuda.jit means we can call it from python by passing cuda\n",
        "    device array objects for \"board\" and \"valid_actions\".\n",
        "\n",
        "    CUDA is structured by 'blocks' and 'threads'. Each block contains multiple\n",
        "    threads. It is the caller's job to say how many of each and their effective\n",
        "    size. For instance, cuda_get_valid_actions[4, (8,8)](...) is the syntax for\n",
        "    \"4 blocks, each of which has 8 'x' and 8 'y' threads\".\n",
        "\n",
        "    We use a separate block per board, and each thread evaluates a single board\n",
        "    coordinate.\n",
        "    \"\"\"\n",
        "    # From numba documentation: const.array_like copies the given array into\n",
        "    # constant GPU memory *at compile time*. This means the copy only happens\n",
        "    # the first time this function is called.\n",
        "    rays = cuda.const.array_like(RAYS)\n",
        "\n",
        "    # We access boards through the 1d block index (boardIdx.x) and positions\n",
        "    # within the board through the thread index x and y.\n",
        "    tx = cuda.threadIdx.x\n",
        "    ty = cuda.threadIdx.y\n",
        "    bx = cuda.blockIdx.x\n",
        "    \n",
        "    # If a space is occupied, it's an illegal move and we're done. Store a zero\n",
        "    # at coordinate (ty, tx) on board bx.\n",
        "    if boards[bx, ty, tx] != 0:\n",
        "        valid_actions[bx, ty, tx] = 0\n",
        "        return\n",
        "\n",
        "    # Try out all 8 ray directions. If any one is a hit, set return value to 1\n",
        "    # and break.\n",
        "    for i in range(8):\n",
        "        hit = cuda_cast_rays(boards[bx], ty, tx, rays[i], player)\n",
        "        if hit:\n",
        "            valid_actions[bx, ty, tx] = 1\n",
        "            return\n",
        "    \n",
        "    # We tried all of the rays and none were legal capturing moves.\n",
        "    valid_actions[bx, ty, tx] = 0\n",
        "\n",
        "@cuda.jit\n",
        "def cuda_step(boards, actions, player):\n",
        "    \"\"\"Helper CUDA kernel for playing actions across multiple boards in parallel.\n",
        "\n",
        "    Boards must be [num_games, n, n], actions must be [num_games, 2] and player\n",
        "    must be 1 or 2. Actions are (row, col) indices of the player's stone.\n",
        "\n",
        "    WARNING: no legality checks are performed here! It is the caller's\n",
        "    responsibility to ensure that all moves are legal.\n",
        "\n",
        "    Boards are indexed by block x and the 8 ray directions by thread x\n",
        "\n",
        "    If either of actions[b,:] is -1, treated as pass.\n",
        "    \"\"\"\n",
        "    # As in cuda_get_valid_actions, this copy happens once at compile time\n",
        "    rays = cuda.const.array_like(RAYS)\n",
        "\n",
        "    # Grab thread and block index from cuda context\n",
        "    tx = cuda.threadIdx.x\n",
        "    bx = cuda.blockIdx.x  \n",
        "\n",
        "    # Unpack the row, col coordinate of the action\n",
        "    act_y, act_x = actions[bx]\n",
        "\n",
        "    # If -1, the player passed and there is nothing to do.\n",
        "    if act_x == -1 or act_y == -1:\n",
        "        return\n",
        "\n",
        "    # Search a different ray direction on each of 8 threads. ASSUMES the move\n",
        "    # is legal.\n",
        "    is_hit = cuda_cast_rays(boards[bx], act_y, act_x, rays[tx], player)\n",
        "    if is_hit:\n",
        "        dy, dx = rays[tx]\n",
        "        opponent = 3 - player\n",
        "        y, x = act_y+dy, act_x+dx\n",
        "        while boards[bx, y, x] == opponent:\n",
        "            # Modify the board, flipping opponent to player\n",
        "            boards[bx, y, x] = player\n",
        "            y, x = y+dy, x+dx\n",
        "    \n",
        "    # Only need to do this once, so just let thread #0 handle it: update\n",
        "    # the board at location of the action\n",
        "    if tx == 0:\n",
        "        boards[bx, act_y, act_x] = player\n",
        "\n",
        "# ====================================\n",
        "# === BEGIN OTHELLOGAME DEFINITION ===\n",
        "# ====================================\n",
        "\n",
        "class OthelloGame(object):\n",
        "    PLAYER1 = 1\n",
        "    PLAYER2 = 2\n",
        "    ACTIVE = 0\n",
        "    OVER = 1\n",
        "    PASS = (-1, -1)\n",
        "    # Unicode characters for dark and light circles if printing to console.\n",
        "    # By convention, player 1 is dark and 2 is light.\n",
        "    GAME_SYMBOLS = ['_', '\\u25cf', '\\u25cb']\n",
        "\n",
        "    def __init__(self, n_games=1, n=8):\n",
        "        \"\"\"Create a new batch of n_games othello games each with board size [n,n]\n",
        "\n",
        "        The game states are managed on the GPU with the help of CUDA kernels.\n",
        "\n",
        "        We expect all n games to update in lockstep, so that there is only one\n",
        "        'current_player' shared by all of them.\n",
        "        \"\"\"\n",
        "        self.n_games = n_games\n",
        "        self.n = n\n",
        "        self.boards = torch.zeros((n_games, n, n), dtype=torch.float32, device='cuda')\n",
        "        # Initial positions: 4 stones in the center alternating color\n",
        "        self.boards[:, [self.n//2, self.n//2-1], [self.n//2, self.n//2-1]] = OthelloGame.PLAYER2\n",
        "        self.boards[:, [self.n//2, self.n//2-1], [self.n//2-1, self.n//2]] = OthelloGame.PLAYER1\n",
        "        # Always start with player 1 (dark stones by convention)\n",
        "        self.current_player = OthelloGame.PLAYER1\n",
        "        # All games are initially active\n",
        "        self.game_status = OthelloGame.ACTIVE*torch.ones(n_games, dtype=torch.int32, device='cuda')\n",
        "        # If both players pass the game is over. Keep track of whether the last\n",
        "        # player passed.\n",
        "        self.last_player_pass = torch.zeros(n_games, dtype=torch.int32, device='cuda')\n",
        "\n",
        "        # _valid_moves is a container for storing which moves are valid. It is\n",
        "        # zero everywhere except where it is legal to play, which is 1.\n",
        "        self._valid_moves = torch.zeros_like(self.boards)\n",
        "\n",
        "        # Create a CUDA device array copy of each GPU tensor. These share GPU\n",
        "        # memory with torch tensors so any changes made in CUDA kernels are seen\n",
        "        # by torch.\n",
        "        self._cuda_boards = torch2cuda(self.boards)\n",
        "        self._cuda_valid = torch2cuda(self._valid_moves)\n",
        "\n",
        "        # Populate valid moves using CUDA-ized algorithm\n",
        "        cuda_get_valid_actions[self.n_games, (self.n, self.n)](self._cuda_boards, self.current_player, self._cuda_valid)\n",
        "\n",
        "        # Initialize 'zobrist hash' table for get_uid() function\n",
        "        self._new_zobrist_table()\n",
        "\n",
        "    def step(self, actions:torch.Tensor):\n",
        "        \"\"\"Place a stone for the current player at location action=(row,col)\n",
        "        separately per board. Expected size of 'actions' is [n_games, 2].\n",
        "\n",
        "        Action coordinates of -1 treated as passing.\n",
        "\n",
        "        Actions are ignored for all games that are not in the ACTIVE state.\n",
        "\n",
        "        Note: no legality checking here! It is the caller's responsibility to\n",
        "        ensure 'action' is legal, e.g. by checking that action is in\n",
        "        game.get_available_actions().\n",
        "        \"\"\"\n",
        "        # Ensure actions is a torch tensor and on the GPU. Ensure its shape is\n",
        "        # [n_games, 2] even if the input was, say, just a (row,col) tuple.\n",
        "        actions = torch.as_tensor(actions, device='cuda').view(-1, 2)\n",
        "\n",
        "        # Ensure actions are \"pass\" for games that have completed. This tells\n",
        "        # the cuda_step kernel to ignore these games.\n",
        "        actions[self.game_status == OthelloGame.OVER, :] = -1\n",
        "\n",
        "        # Update board state with call to CUDA kernel cuda_step.\n",
        "        cuda_step[self.n_games, RAYS.shape[0]](self._cuda_boards, torch2cuda(actions), self.current_player)\n",
        "\n",
        "        # Advance to the next player (2->1 and 1->2).\n",
        "        self.current_player = 3-self.current_player\n",
        "\n",
        "        # Count passes - any games with 2 passes in a row is flagged as being over.\n",
        "        is_pass = torch.any(actions == -1, dim=1)\n",
        "        self.last_player_pass[~is_pass] = 0\n",
        "        self.last_player_pass[is_pass] += 1\n",
        "\n",
        "        # Flag games that have finished\n",
        "        is_game_over = self.last_player_pass >= 2\n",
        "        self.game_status[is_game_over] = OthelloGame.OVER\n",
        "\n",
        "        # Look for valid moves in the updated boards for the new player\n",
        "        cuda_get_valid_actions[self.n_games, (self.n, self.n)](self._cuda_boards, self.current_player, self._cuda_valid)\n",
        "\n",
        "    def copy_state(self, idx=None):\n",
        "        # Grab the minimal amount of state info to be able to return the board\n",
        "        # to its current configuration later with a call to paste_state.\n",
        "        # NOTE: this copy is single-use only; once it is pasted, it will be\n",
        "        # modified in-place!\n",
        "        return {\"boards\": self.boards.clone(),\n",
        "                \"game_status\": self.game_status.clone(),\n",
        "                \"last_player_pass\": self.last_player_pass.clone(),\n",
        "                \"current_player\": self.current_player}\n",
        "\n",
        "\n",
        "    def paste_state(self, state):\n",
        "        self.__dict__.update(state)\n",
        "        self._cuda_boards = torch2cuda(self.boards)\n",
        "        self._cuda_valid = torch2cuda(self._valid_moves)\n",
        "        cuda_get_valid_actions[self.n_games, (self.n, self.n)](self._cuda_boards, self.current_player, self._cuda_valid)\n",
        "\n",
        "    def render(self, mode='human'):\n",
        "        boards = self.boards.cpu()\n",
        "        acts = self.get_available_actions()\n",
        "        for g in range(self.n_games):\n",
        "            if mode == 'text':\n",
        "                print(\"=\"*(self.n-1) + (f\"{OthelloGame.GAME_SYMBOLS[self.current_player]}\")*2 + \"=\"*(self.n-1))\n",
        "                for row in boards[g]:\n",
        "                    print(\" \".join(OthelloGame.GAME_SYMBOLS[v] for v in row))\n",
        "                print(\"=\"*(self.n*2))\n",
        "            else:\n",
        "                # Plot with matplotlib patches. Player 1 is dark and 2 is light.\n",
        "                fig, ax = plt.subplots(figsize=(6,6))\n",
        "                ax.set_aspect('equal')\n",
        "                rect = patches.Rectangle((0,0),self.n,self.n,linewidth=1,edgecolor='k',facecolor='g')\n",
        "                ax.add_patch(rect)\n",
        "                ax.set_xlim(0,self.n)\n",
        "                ax.set_ylim(0,self.n)\n",
        "                ax.set_yticklabels([])\n",
        "                ax.set_xticklabels([])\n",
        "                ax.tick_params(length=0)\n",
        "                ax.grid(which='major', zorder=0, c='k')\n",
        "                for i in range(self.n):\n",
        "                    for j in range(self.n):\n",
        "                        if boards[g,i,j]>0:\n",
        "                            # Draw stones\n",
        "                            c = (0,0,0,1) if boards[g,i,j]==OthelloGame.PLAYER1 else (1,1,1,1)\n",
        "                            circ = patches.Circle((j+.525,self.n-1-i+.5), .4, \n",
        "                                                facecolor=c, linewidth=1,edgecolor='k')\n",
        "                            ax.add_patch(circ)\n",
        "                        elif (i,j) in acts[g]:\n",
        "                            # Draw transluscent stones on legal move positions\n",
        "                            c = (0,0,0,0.25) if self.current_player==OthelloGame.PLAYER1 else (1,1,1,0.25)\n",
        "                            circ = patches.Circle((j+.525,self.n-1-i+.5), .4, facecolor=c, linewidth=0)\n",
        "                            ax.add_patch(circ)\n",
        "                plt.show()\n",
        "\n",
        "    def get_available_actions(self):\n",
        "        \"\"\"Get a list of all legal actions for the current player (passing not included).\n",
        "        \n",
        "        Note: assumes self._valid_moves is up to date!\n",
        "        \"\"\"\n",
        "        acts = [None]*self.n_games\n",
        "        for g in range(self.n_games):\n",
        "            i, j = torch.where(self._valid_moves[g] == 1)\n",
        "            acts[g] = list(zip(i.cpu().numpy(), j.cpu().numpy()))\n",
        "        return acts\n",
        "\n",
        "    def are_games_over(self):\n",
        "        return torch.all(self.game_status == OthelloGame.OVER).item()\n",
        "\n",
        "    def score_games(self):\n",
        "        \"\"\"Compute [n_games, 2] tuple containing \"scores\" for player 1 and 2 in\n",
        "        each column. Assumes game is over.\n",
        "\n",
        "        scores[:,0] is +1 if player 1 won, 0 for draw, or -1 for loss\n",
        "        scores[:,1] is the same from player 2's perspective\n",
        "        \"\"\"\n",
        "        flat_boards = self.boards.view(self.n_games, -1)\n",
        "        n_dark = torch.sum(flat_boards == OthelloGame.PLAYER1, dim=1)\n",
        "        n_lite = torch.sum(flat_boards == OthelloGame.PLAYER2, dim=1)\n",
        "        scores = torch.zeros_like(n_dark)\n",
        "        scores[n_dark > n_lite] = 1\n",
        "        scores[n_dark < n_lite] = -1\n",
        "        return torch.stack([scores,-scores], dim=1)\n",
        "    \n",
        "    def get_uid(self, out=None):\n",
        "        \"\"\"Return highly-probably-unique identifier for each game. Output is a\n",
        "        tensor of [n_games] int64 values.\n",
        "\n",
        "        Algorithm is Zobrist hashing.\n",
        "        \"\"\"\n",
        "        flat_boards = self.boards.reshape(self.n_games, -1)\n",
        "        # Board values are 0,1,2 for each player. Use these as an index into\n",
        "        # the self._hash_bits table: each state [0,1,2] at each position in the\n",
        "        # [n,n] board has a unique random bit string. The Zobrist algorithm is\n",
        "        # simply an XOR of all of these values.\n",
        "        bit_strings = self._hash_bits[torch.arange(self.n**2), flat_boards.long()]\n",
        "        # Manually xor everything together (there is no pytorch operation for\n",
        "        # xor along a dimension, unfortunately)\n",
        "        if out is None:\n",
        "            out = torch.zeros(self.n_games, dtype=torch.int64, device='cuda')\n",
        "        for i in range(self.n**2):\n",
        "            torch.bitwise_xor(out, bit_strings[:,i], out=out)\n",
        "        return out\n",
        "\n",
        "    def _new_zobrist_table(self):\n",
        "        self._hash_bits = torch.randint(2**63-1, size=(self.n*self.n, 3), dtype=torch.int64, device='cuda')\n",
        "        # Sanity check that no 2 random strings were identical\n",
        "        assert len(self._hash_bits.flatten().unique()) == self.n*self.n*3\n",
        "\n",
        "################################\n",
        "#### INTERACTIVE GAME BOARD ####\n",
        "################################\n",
        "\n",
        "def temporary_info(message, clear=False):\n",
        "    \"\"\"Output overwritable message to the console. If clear=True, overwrite all old messages\n",
        "    \"\"\"\n",
        "    if clear:\n",
        "        output.clear(output_tags=\"temporary-info\")\n",
        "    with output.use_tags(\"temporary-info\"):\n",
        "        print(message)\n",
        "\n",
        "class InteractiveOthelloGame(object):\n",
        "    def __init__(self, game=None, player1=\"human\", player2=\"human\", n=8):\n",
        "        if game is None:\n",
        "            self.game = OthelloGame(n_games=1, n=n)\n",
        "        else:\n",
        "            self.game = game\n",
        "            assert game.n_games == 1, \"Can't handle >1 games interactively!\"\n",
        "        self.players = [player1, player2]\n",
        "\n",
        "    def next_turn(self):\n",
        "        self.available_acts = self.game.get_available_actions()[0]\n",
        "        if self.game.game_status == OthelloGame.ACTIVE:\n",
        "            the_player = self.players[self.game.current_player-1]\n",
        "            temporary_info(f\"Begin Player {self.game.current_player}'s turn\", clear=True)\n",
        "            if the_player == \"human\":\n",
        "                if len(self.available_acts) > 0:\n",
        "                    temporary_info(\"Input move by clicking the board\")\n",
        "                else:\n",
        "                    temporary_info(\"No legal moves! Click anywhere on the board to pass.\")\n",
        "                self.redraw()\n",
        "                self._reregister_click_callback()\n",
        "            else:\n",
        "                temporary_info(\"Waiting for AI to complete...\")\n",
        "                ai_action = the_player.select_move(self.game)\n",
        "                self.game.step(ai_action)\n",
        "                self.next_turn()\n",
        "        else:\n",
        "            self.redraw()\n",
        "            temporary_info(\"GAME OVER!\", clear=True)\n",
        "            values = self.game.score_games()[0]\n",
        "            if values[0] == 0:\n",
        "                print(f\"Game ended in a draw\")\n",
        "            else:\n",
        "                winner_id = 1+np.argmax(values)\n",
        "                try:\n",
        "                    name = self.players[winner_id-1].name\n",
        "                except AttributeError:\n",
        "                    name = f\"Player {winner_id}\"\n",
        "                print(f\"{name} is the winner!\")\n",
        "\n",
        "    def handle_user_click(self, x_pix, y_pix, plot_width, plot_height):\n",
        "        if len(self.available_acts) == 0:\n",
        "            # If there is not available action all you can do is pass\n",
        "            self.game.step(OthelloGame.PASS)\n",
        "            self.next_turn()\n",
        "        else:\n",
        "            cell_width = plot_width / self.game.boards.shape[2]\n",
        "            cell_height = plot_height / self.game.boards.shape[1]\n",
        "            cell_x, cell_y = int(x_pix / cell_width), int(y_pix / cell_height)\n",
        "            action = (cell_y, cell_x)\n",
        "            if action in self.available_acts:\n",
        "                self.game.step(action)\n",
        "                self.next_turn()\n",
        "            else:\n",
        "                temporary_info(f\"Available actions are {self.available_acts} but you clicked {action}\")\n",
        "\n",
        "    def _reregister_click_callback(self):\n",
        "        # Inject javascript which will detect a click and invoke a function called 'pass_to_python_handler'\n",
        "        display(IPython.display.Javascript(\"\"\"\n",
        "        var plot_element = document.querySelector(\".output_image\").firstElementChild;\n",
        "        plot_element.onclick = function(event){\n",
        "            google.colab.kernel.invokeFunction(\"pass_to_python_handler\", [event.offsetX, event.offsetY, plot_element.width, plot_element.height], {});\n",
        "        };\n",
        "        \"\"\"))\n",
        "\n",
        "        # Tell colab that when 'pass_to_python_handler' is called in JS, it should\n",
        "        # call self.handle_user_click in python\n",
        "        output.register_callback(\"pass_to_python_handler\", self.handle_user_click)\n",
        "\n",
        "    def redraw(self):\n",
        "        # Clear previous output\n",
        "        output.clear(output_tags='othello-interactive')\n",
        "        # Draw a fresh plot and store the figsize in pixels\n",
        "        with output.use_tags('othello-interactive'):\n",
        "            self.game.render()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lZ0g92gbX1kZ"
      },
      "source": [
        "A `game` is a class instance of `OthelloGame`. This environment is easy to use and handles all the game mechanics internally.\n",
        "\n",
        "For speed reasons later on, a single `OthelloGame` instance can be used to manage an entire collection of games being played at once with the `n_games` argument, and expensive board-management is done on the GPU. For now, don't worry about any of these gory implementation details. We'll look at playing one game at a time.\n",
        "\n",
        "Let's look at the initial board layout: Player 1 stones are black, Player 2 stones are white, and the transparent stones indicate legal moves for the current player."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ant8wkFLctxj"
      },
      "source": [
        "game = OthelloGame()\n",
        "game.render()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y0gyQsdVY063"
      },
      "source": [
        "`.get_available_actions()` will tell you which board positions are available for play from the perspective of the current player (`board.current_player`).  \n",
        "**Note**: `get_available_actions()` returns a list of lists of (row, col) tuples. One list per game being played in parallel. Here, we get the 0th list, which corresponds to the 0th gameboard. Later we will play many games in parallel."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Q_L5dcT1Ykl_"
      },
      "source": [
        "# the 0th list, which corresponds to the 0th gameboard.\n",
        "actions = game.get_available_actions()[0]\n",
        "print(\"Available actions: \", actions)\n",
        "print(\"Current player: \", game.current_player)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2RRBAAqUZcUF"
      },
      "source": [
        "`.step(action)` will play a move, flip the correct pieces, and switch the current player."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qIV2Jug8CZfZ"
      },
      "source": [
        "game = OthelloGame()\n",
        "game.step(actions[0])\n",
        "game.render()\n",
        "print(\"Available actions: \", game.get_available_actions()[0])\n",
        "print(\"Current player: \", game.current_player)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IVWgw4aqZ217"
      },
      "source": [
        "The game ends when both players are forced to pass in a row. Passing is only considered legal if a player has no other moves available.\n",
        "\n",
        "`game.step(action)` doesn't enforce that moves are legal - use this power wisely!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eX0Co50SYuwc"
      },
      "source": [
        "print(f\"PASS by playing {OthelloGame.PASS}\")\n",
        "game = OthelloGame()\n",
        "print(game.are_games_over())\n",
        "game.step(OthelloGame.PASS) # Player 1 pass\n",
        "print(game.are_games_over())\n",
        "game.step(OthelloGame.PASS) # Player 2 pass\n",
        "print(game.are_games_over())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bQI2UxEvj3Dn"
      },
      "source": [
        "## Section 1.2: Playing multiple games at once\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "c7_TYi1SJCXL",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Using More Data\n",
        "video = YouTubeVideo(id=\"zLEQ0xQDK8o\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hjt6Zt0ITyrc"
      },
      "source": [
        "Deep learning almost universally works with *batches* of examples, rather than one example at a time. The biggest reason is that one example is not representative of the ensemble of data; we want to learn general aspects, not memorize that example. The other reason is that GPUs are phenomenal at parallel evaluation. We're going to start thinking like this from day 1.\n",
        "\n",
        "Instead of 1 game at a time, we build a simulator that plays `n` games at a time. We can initialize two games with `OthelloGame(n_games=2)`:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WltjDpZrj3Vk"
      },
      "source": [
        "game = OthelloGame(n_games=2)\n",
        "game.render()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0RiG5RCJc72K"
      },
      "source": [
        "`.get_available_actions()` now returns a list with one entry for each game."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CDnFUX2yc79V"
      },
      "source": [
        "actions = game.get_available_actions()\n",
        "print(\"Available actions: \", actions)\n",
        "print(\"Current player: \", game.current_player)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7EcZh6CUc8Dy"
      },
      "source": [
        "Note that there is still just one current player. This is because we play both games in parallel.\n",
        "\n",
        "`.step()` will now take a *matrix* of actions, with one row per game. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-WrK42yLc8MU"
      },
      "source": [
        "game = OthelloGame(n_games = 2)\n",
        "actions = np.array([[2,3],    # game 1\n",
        "                    [3,2]])   # game 2\n",
        "game.step(actions)\n",
        "game.render()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UHd-0VV0j3mM"
      },
      "source": [
        "## Section 1.3: The super-basics of Pytorch Tensors\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PrgDLbEyT5H7",
        "cellView": "form"
      },
      "source": [
        "#@title Video: A Whirlwind Tour of Pytorch\n",
        "video = YouTubeVideo(id=\"U5Bdnt7aNvA\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RNGkpfIpT7DR"
      },
      "source": [
        "\n",
        "This class will use Pytorch. Next week will be our real introduction to Pytorch – so hold on! This will seem like a lot if you haven't learned this language before.\n",
        "\n",
        "Today, all that is needed to know is that Pytorch acts a lot like Numpy. We can create Tensors (instead of arrays):"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pECPevHZj4Rr"
      },
      "source": [
        "t = torch.Tensor([[3,2],[4,2]])\n",
        "print(t)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mcCnne39gHKD"
      },
      "source": [
        "Run torch commands on these tensors:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zBdR0nU6gHcb"
      },
      "source": [
        "print(torch.sum(t, dim=1))\n",
        "print(torch.argmax(t, dim=0))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Vw8_OZbmgVMK"
      },
      "source": [
        "And do essentially anything we would do in Numpy:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DYDMNWi8gXJg"
      },
      "source": [
        "torch.rand(2,1,1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C-IM5CugghBK"
      },
      "source": [
        "Many more details will come tomorrow. There will always be more to Torch than we can tell you ourselves. The documentation is your friend!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wYBatwZgpVb6"
      },
      "source": [
        "## Exercise 0: a Pytorch warm-up\n",
        "As a test if your knowledge, let's see if you can use some Pytorch skills to learn something about the way we set up Othello boards.\n",
        "\n",
        "Code something below to determine how many piece of a given player are on each board.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2a5_Hsm8pnxx"
      },
      "source": [
        "def count_pieces(game:OthelloGame, whoami):\n",
        "    \"\"\"A function to sum the number of pieces on the board of a given player\n",
        "    current player (whoami).\n",
        "    Inputs: game, an instance of OthelloGame\n",
        "            whoami (int) 1 or 2, whose turn it is\n",
        "    Outputs:\n",
        "            populous_row: a 1d pytorch tensor of shape (n_games)\"\"\"\n",
        "\n",
        "    #####################################################################\n",
        "    # Fill in missing code (...),\n",
        "    # then remove or comment the line below to test your function.\n",
        "    ## hint 1: inspect game.boards\n",
        "    ## hint 2: which dimensions should you sum over?\n",
        "    # IMPORTANT: Make sure to use PyTorch tensors and NOT numpy arrays!\n",
        "    raise NotImplementedError(\"Complete the count_pieces function\")\n",
        "    #####################################################################\n",
        "\n",
        "    n_my_pieces = ...\n",
        "\n",
        "    return n_my_pieces\n",
        "\n",
        "## Uncomment to test code\n",
        "# game = OthelloGame(n_games=2) \n",
        "# for turn in range(4): # take some actions in both games\n",
        "#     actions = actions_as_tensor(game.get_available_actions())\n",
        "#     act = actions[[0,1],[0,1]]\n",
        "#     game.step(act)\n",
        "\n",
        "# game.render()\n",
        "# print(count_pieces(game, 1))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AvZ9B9n_WwJD"
      },
      "source": [
        "# to_remove solution\n",
        "def count_pieces(game:OthelloGame, whoami):\n",
        "    \"\"\"A function to sum the number of pieces on the board of a given player\n",
        "    current player (whoami).\n",
        "    Inputs: game, an instance of OthelloGame\n",
        "            whoami (int) 1 or 2, whose turn it is\n",
        "    Outputs:\n",
        "            populous_row: a 1d pytorch tensor of shape (n_games)\"\"\"\n",
        "\n",
        "    n_my_pieces = torch.sum(game.boards==whoami, dim=(1,2))\n",
        "\n",
        "    return n_my_pieces\n",
        "\n",
        "## Uncomment to test code\n",
        "game = OthelloGame(n_games=2) \n",
        "for turn in range(4): # take some actions in both games\n",
        "    actions = actions_as_tensor(game.get_available_actions())\n",
        "    act = actions[[0,1],[0,1]]\n",
        "    game.step(act)\n",
        "\n",
        "game.render()\n",
        "print(count_pieces(game, 1))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vdPrhOBC_Wgz"
      },
      "source": [
        "---\n",
        "# Section 2: Thinking harder about gameplaying"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SNeMermGI7-c",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Gaining Intuition for how to use DL\n",
        "video = YouTubeVideo(id=\"oTLO7MY67Kc\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "try:t2;\n",
        "except NameError: t2=time.time()\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FL-Ae7wUaWeh",
        "cellView": "form"
      },
      "source": [
        "#@title Create an interactive point-and-click interface (just run me)\n",
        "\n",
        "def temporary_info(message, clear=False):\n",
        "    \"\"\"Output overwritable message to the console. If clear=True, overwrite all old messages\n",
        "    \"\"\"\n",
        "    if clear:\n",
        "        output.clear(output_tags=\"temporary-info\")\n",
        "    with output.use_tags(\"temporary-info\"):\n",
        "        print(message)\n",
        "\n",
        "class InteractiveOthelloGame(object):\n",
        "    def __init__(self, game=None, player1=\"human\", player2=\"human\", n=8):\n",
        "        if game is None:\n",
        "            self.game = OthelloGame(n_games=1, n=n)\n",
        "        else:\n",
        "            self.game = game\n",
        "            assert game.n_games == 1, \"Can't handle >1 games interactively!\"\n",
        "        self.players = [player1, player2]\n",
        "\n",
        "    def next_turn(self):\n",
        "        self.available_acts = self.game.get_available_actions()[0]\n",
        "        if self.game.game_status == OthelloGame.ACTIVE:\n",
        "            the_player = self.players[self.game.current_player-1]\n",
        "            temporary_info(f\"Begin Player {self.game.current_player}'s turn\", clear=True)\n",
        "            if the_player == \"human\":\n",
        "                if len(self.available_acts) > 0:\n",
        "                    temporary_info(\"Input move by clicking the board\")\n",
        "                else:\n",
        "                    temporary_info(\"No legal moves! Click anywhere on the board to pass.\")\n",
        "                self.redraw()\n",
        "                self._reregister_click_callback()\n",
        "            else:\n",
        "                temporary_info(\"Waiting for AI to complete...\")\n",
        "                ai_action = the_player.select_move(self.game)\n",
        "                self.game.step(ai_action)\n",
        "                self.next_turn()\n",
        "        else:\n",
        "            self.redraw()\n",
        "            temporary_info(\"GAME OVER!\", clear=True)\n",
        "            values = self.game.score_games()\n",
        "            if values[0] == 0:\n",
        "                temporary_info(f\"Game ended in a draw\")\n",
        "            else:\n",
        "                winner_id = 1+np.argmax(values)\n",
        "                try:\n",
        "                    name = self.players[winner_id-1].name\n",
        "                except AttributeError:\n",
        "                    name = f\"Player {winner_id}\"\n",
        "                temporary_info(f\"{name} is the winner!\")\n",
        "\n",
        "    def handle_user_click(self, x_pix, y_pix, plot_width, plot_height):\n",
        "        if len(self.available_acts) == 0:\n",
        "            # If there is not available action all you can do is pass\n",
        "            self.game.step(OthelloGame.PASS)\n",
        "            self.next_turn()\n",
        "        else:\n",
        "            cell_width = plot_width / self.game.boards.shape[2]\n",
        "            cell_height = plot_height / self.game.boards.shape[1]\n",
        "            cell_x, cell_y = int(x_pix / cell_width), int(y_pix / cell_height)\n",
        "            action = (cell_y, cell_x)\n",
        "            if action in self.available_acts:\n",
        "                self.game.step(action)\n",
        "                self.next_turn()\n",
        "            else:\n",
        "                temporary_info(f\"Available actions are {self.available_acts} but you clicked {action}\")\n",
        "\n",
        "    def _reregister_click_callback(self):\n",
        "        # Inject javascript which will detect a click and invoke a function called 'pass_to_python_handler'\n",
        "        display(IPython.display.Javascript(\"\"\"\n",
        "        var plot_element = document.querySelector(\".output_image\").firstElementChild;\n",
        "        plot_element.onclick = function(event){\n",
        "            google.colab.kernel.invokeFunction(\"pass_to_python_handler\", [event.offsetX, event.offsetY, plot_element.width, plot_element.height], {});\n",
        "        };\n",
        "        \"\"\"))\n",
        "\n",
        "        # Tell colab that when 'pass_to_python_handler' is called in JS, it should\n",
        "        # call self.handle_user_click in python\n",
        "        output.register_callback(\"pass_to_python_handler\", self.handle_user_click)\n",
        "\n",
        "    def redraw(self):\n",
        "        # Clear previous output\n",
        "        output.clear(output_tags='othello-interactive')\n",
        "        # Draw a fresh plot and store the figsize in pixels\n",
        "        with output.use_tags('othello-interactive'):\n",
        "            self.game.render()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pl617RtbY8vX"
      },
      "source": [
        "## Exercise 1: Play a game\n",
        "In breakout rooms, play a quick game against another student."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "87KWEBsWilfo"
      },
      "source": [
        "interface = InteractiveOthelloGame(n=6)\n",
        "interface.next_turn()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "obBUwuk1C4xj"
      },
      "source": [
        "Describe how you were thinking during the game. Were you optimizing something? Did you do some kind of planning?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TP2-IpjO-jl-",
        "cellView": "form"
      },
      "source": [
        "play_analysis= 'something' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zVQWgg_wLa2r"
      },
      "source": [
        "Run the above cell when you finish writing."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MLX_gZ8b6_Si"
      },
      "source": [
        "---\n",
        "# Section 3: Planning and value judgement: two complementary strategies for game playing\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RNdGNG7y_WY3",
        "cellView": "form"
      },
      "source": [
        "#@title Video: What Pieces Are Needed for Teaching a Computer Othello\n",
        "\n",
        "param_list = ['play_analysis']\n",
        "for param in param_list:\n",
        "    if param not in locals():\n",
        "        raise NameError(\"Please make sure to run the cell \"\n",
        "                        \"after submitting your answer for <{}>!\".format(param))\n",
        "\n",
        "video = YouTubeVideo(id=\"f6u-lq6tC-k\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "try: t3;\n",
        "except NameError: t3=time.time()\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fP0Qsb_oVgay"
      },
      "source": [
        "\n",
        "The magic of deep learning requires intelligently formulating your objective function. Gameplaying can be reduced to **two** separate objectives, each of which we will eventually train a neural network for. These objectives are **value** and **policy**. Very informally,\n",
        "\n",
        "  - **Value** is a snap-judgement estimation of how likely a current game state will result in a final win (under some implicit game-playing strategy).\n",
        "  - **Policy** is also a snap judgement that directly outputs the next promising move given a board state."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eKx3sxVHIQlH"
      },
      "source": [
        "## Section 3.1: More about value\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N0AeceAxVkOx"
      },
      "source": [
        "\n",
        "Let's think more about value. A value estimate is the probablility that a given board state $s$ will result in a win.\n",
        "\n",
        "The value of a winning board state $V(s_W)$ is obviously 1. What about a loss? Typically $V(s_L)=0$, but today we're going to rescale $V(s)$ so that a loss has value $V(s_L)=-1$. Draws have $V(s_D)=0$.\n",
        "\n",
        "It is crucial to understand that a value estimate board depends on an assumption of how moves are played by you and your opponents (your \"policies\"). A board state progresses to a terminal board state through these actions!\n",
        "\n",
        "You may recall that value can be recursively calculated via Bellman's equation and knowledge of the policy. That's *not* how we're going to view it today. Instead, we'll imagine $V(s)$ to be calculated by neural intuition: a snap-judgement estimate of whether player 1 will win, lose, or draw."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L-0XOKXDilm4"
      },
      "source": [
        "\n",
        "## Exercise 2: Write your own value function\n",
        "\n",
        "Warning to TAs: do not spend more than 20 minutes on this. If your timer hits 20 minutes, provide the solution."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8uKd_Qw-8tIv"
      },
      "source": [
        "def my_board_value(board_states, whoami=1):\n",
        "    \"\"\"\n",
        "    Returns the *value* of a board or set of boards.\n",
        "    \n",
        "    The value is a number between -1 and 1 that indicates how good that board is\n",
        "    for the player. Specifically, the value is a measure of confidence\n",
        "    that a board state will lead to a final win.\n",
        "    1 is an estimated win, -1 a loss, and 0 means a draw.\n",
        "\n",
        "    Inputs: board_states (torch tensor with size [n_games, n, n]):\n",
        "                0s are positions with no piece:\n",
        "                1 is player 1, \n",
        "                2 is player 2.\n",
        "            whoami is an integer (1 or 2) indicating which player \"you\" are.\n",
        "    Return +1 if the \"whoami\" player is sure to win, -1 if the \"other\" player is\n",
        "    sure to win, 0 for a draw, or any value in between.\n",
        "\n",
        "    Outputs: torch array (1d tensor) of values, one per n_games \n",
        "\n",
        "    Note: this function must not modify board_states!\n",
        "    \"\"\"\n",
        "    n_games, n, _ = board_states.size()\n",
        "    #####################################################################\n",
        "    # Fill in missing code (...),\n",
        "    # then remove or comment the line below to test your function.\n",
        "    # Your heuristic doesn't need to be perfect. Just try to capture something\n",
        "    # about \"goodness\" or \"badness\" of a board.\n",
        "    # Remember that board_states is 3d (n_games, x, y) \n",
        "    # and we want a value for each game.\n",
        "    # You may need to code multiple lines and use auxiliary variables\n",
        "    # IMPORTANT: Make sure to use PyTorch tensors and NOT numpy arrays!\n",
        "    raise NotImplementedError(\"Complete the my_board_value function\")\n",
        "    #####################################################################\n",
        "    values = ... \n",
        "\n",
        "    assert type(values) == torch.Tensor, \"type(values) must be torch.Tensor\" \n",
        "    return values if whoami == 1 else -values\n",
        "\n",
        "game = OthelloGame()\n",
        "game.step(game.get_available_actions()[0][0])\n",
        "game.step(game.get_available_actions()[0][1])\n",
        "game.step(game.get_available_actions()[0][0])\n",
        "game.render()\n",
        "\n",
        "## uncomment the lines below to test your function\n",
        "#print(f\"Value for player 1 (dark) is {my_board_value(game.boards, whoami=1)}\")\n",
        "#print(f\"Value for player 2 (light) is {my_board_value(game.boards, whoami=2)}\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9zI4btk98tLV"
      },
      "source": [
        "# to_remove solution\n",
        "def my_board_value(board_states, whoami=1):\n",
        "    \"\"\"\n",
        "    Returns the *value* of a board or set of boards.\n",
        "    \n",
        "    The value is a number between -1 and 1 that indicates how good that board is\n",
        "    for the player. Specifically, the value is a measure of confidence\n",
        "    that a board state will lead to a final win.\n",
        "    1 is an estimated win, -1 a loss, and 0 means a draw.\n",
        "\n",
        "    Inputs: board_states (torch tensor with size [n_games, n, n]):\n",
        "                0s are positions with no piece:\n",
        "                1 is player 1, \n",
        "                2 is player 2.\n",
        "            whoami is an integer (1 or 2) indicating which player \"you\" are.\n",
        "    Return +1 if the \"whoami\" player is sure to win, -1 if the \"other\" player is\n",
        "    sure to win, 0 for a draw, or any value in between.\n",
        "\n",
        "    Outputs: torch array (1d tensor) of values, one per n_games \n",
        "\n",
        "    Note: this function must not modify board_states!\n",
        "    \"\"\"\n",
        "    n_games, n, _ = board_states.size()\n",
        "    \n",
        "    values = torch.zeros(n_games)\n",
        "    n1 = (board_states == OthelloGame.PLAYER1).view(n_games, -1).sum(dim=1)\n",
        "    n2 = (board_states == OthelloGame.PLAYER2).view(n_games, -1).sum(dim=1)\n",
        "\n",
        "    frac1 = n1/(n1+n2)\n",
        "    values = frac1*2-1\n",
        "\n",
        "    return values if whoami == 1 else -values\n",
        "\n",
        "game = OthelloGame()\n",
        "game.step(game.get_available_actions()[0][0])\n",
        "game.step(game.get_available_actions()[0][1])\n",
        "game.step(game.get_available_actions()[0][0])\n",
        "game.render()\n",
        "\n",
        "## uncomment the lines below to test your function\n",
        "print(f\"Value for player 1 (dark) is {my_board_value(game.boards, whoami=1)}\")\n",
        "print(f\"Value for player 2 (light) is {my_board_value(game.boards, whoami=2)}\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l4sma75n5Wbe"
      },
      "source": [
        "---\n",
        "# Section 4: Training a neural network to approximate value"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tiRGUafL6_uK",
        "cellView": "form"
      },
      "source": [
        "#@title Video: From Heuristics to Neural Values\n",
        "\n",
        "video = YouTubeVideo(id=\"4bOP7KRlPf4\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "try: t4;\n",
        "except NameError: t4=time.time()\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IkFsj2Z-ShCf"
      },
      "source": [
        "Let's build a better value estimator using neural networks. \n",
        "\n",
        "We can only use a neural network if we can formulate an appropriate objective function to train towards, and if we can also build the proper labeled training set. How can we build the training set for a value net?\n",
        "\n",
        "This is trickier than it may seem. A value function always implies the probability you will win *under your typical style of gameplay* and also *under your opponent's style of gameplay*. Value is not a property of the board state alone. \n",
        "\n",
        "## Exercise 3\n",
        "Take 2 minutes in silence (before group discussion) and write down a strategy for training a neural network."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "i805GsUBLgZl",
        "cellView": "form"
      },
      "source": [
        "how_to_train_value= '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GZ58WfbPHcym"
      },
      "source": [
        "Now read out your strategies in your pod."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ideyWMEfKRKP"
      },
      "source": [
        "---\n",
        "## Section 4.1: Using value estimates to play games\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Uc5CvZXYQHBs",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Epsilon Greedy Policy\n",
        "\n",
        "param_list = ['how_to_train_value']\n",
        "for param in param_list:\n",
        "    if param not in locals():\n",
        "        raise NameError(\"Please make sure to run the cell \"\n",
        "                        \"after submitting your answer for <{}>!\".format(param))\n",
        "\n",
        "video = YouTubeVideo(id=\"E7smDLlbjqs\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aDwCaxAGQGN4"
      },
      "source": [
        "As your group probably observed, most (but not all!) strategies for training a value function involve observing games. We need to observe board states, observe the resulting play and outcome, and train a network to predict the outcome.\n",
        "\n",
        "How can we build a gameplaying agent? Luckily, one can use a value function alone to play games. This means we can \"pick ourselves up by our bootstraps\" and use a value function to train a better value function.\n",
        "\n",
        "One strategy to play with a value estimate alone is to play **\"$\\epsilon$-greedy\"** w/r/t your value estimate.\n",
        "$\\epsilon$-greedy playing proceeds by:\n",
        " - choosing a random action with probability $\\epsilon$\n",
        " - with probability $1-\\epsilon$, choose the action that results in a board state with the highest estimated value.\n",
        "\n",
        "### Implement an $\\epsilon$ greedy agent\n",
        "Here we will create a $\\epsilon$-greedy AI that uses your hand-built value function. \n",
        "\n",
        "First, let us introduce the base class `ArtificialPlayer` that will be subclassed in all AI players below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ie_-f3MacDdK"
      },
      "source": [
        "class ArtificialPlayer(object):\n",
        "    \"\"\"A base class for agents that play Othello. Sub-classes must provide a\n",
        "    select_move method which takes in an OthelloGame and returns an action\n",
        "    (in other words, returns a tuple (row,col) to place a stone or None to pass)\n",
        "    \"\"\"\n",
        "    def __init__(self, name=\"AI Agent\"):\n",
        "        self.name = name\n",
        "\n",
        "    def select_move(self, game:OthelloGame):\n",
        "        raise NotImplementedError(\"Must subclass with a playing rule\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "puUbcgY7cMQY"
      },
      "source": [
        "Now, we will create an `EpsilonGreedyAgent` that subclasses `ArtificialPlayer`. By default, this has an `evaluate` function that is simply your `my_board_value` from above. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Qrh0hsRgb8SS"
      },
      "source": [
        "class EpsilonGreedyAgent(ArtificialPlayer):\n",
        "    def __init__(self, epsilon=0.05, name=None):\n",
        "        super(EpsilonGreedyAgent,\n",
        "              self).__init__(name if name is not None else f\"Epsilon={epsilon}\")\n",
        "        self.epsilon = epsilon\n",
        "\n",
        "    def select_move(self, game:OthelloGame):\n",
        "        # Evaluate all actions in all games.\n",
        "        all_values, actions_tensor, actions_list = self.batch_evaluate(game,\n",
        "                                                        game.current_player)\n",
        "        # Initially select the best action in each game.\n",
        "        chosen_action = torch.zeros(game.n_games, 2, dtype=torch.int32)\n",
        "        for g in range(game.n_games):\n",
        "            # If there are no moves available, pass\n",
        "            if len(actions_list[g]) == 0:\n",
        "                chosen_action[g, :] = -1\n",
        "                continue\n",
        "            if np.random.rand() < self.epsilon:\n",
        "                acts = actions_list[g]\n",
        "                random_choice = acts[np.random.randint(len(acts))]\n",
        "                chosen_action[g, :] = torch.as_tensor(random_choice,\n",
        "                                                      dtype=torch.int32)\n",
        "            else:\n",
        "                idx = torch.argmax(all_values[g])\n",
        "                chosen_action[g, :] = actions_tensor[g, idx, :]\n",
        "\n",
        "        return chosen_action\n",
        "\n",
        "    def evaluate(self, game:OthelloGame, whoami:int):\n",
        "        \"\"\"\n",
        "        Evaluation function used in \"select_move\" below.\n",
        "        Subclasses may override evaluate() to create epsilon-greedy players\n",
        "        with a different value function.\n",
        "        \"\"\"\n",
        "        # Default to the hand-coded heuristic from above\n",
        "        return my_board_value(game.boards, whoami)\n",
        "    \n",
        "    def batch_evaluate(self, game:OthelloGame, whoami:int):\n",
        "        \"\"\"Call evaluate() once per available action. Returns three things:\n",
        "            - 'values' a [n_games, n_actions] tensor containing values estimated\n",
        "              for each action per game.\n",
        "            - 'actions_tensor' a [n_games, n_actions, 2] tensor containing the\n",
        "              (row,col) coordinates of corresponding actions.\n",
        "            - 'actions_list' the output of game.get_available_actions\n",
        "\n",
        "        Handles the case of game.n_games>1 somewhat efficiently. If different\n",
        "        games have different #s of actions, then n_actions is the maximum number\n",
        "        of actions in any of the games. 'actions' is padded with -1 on games\n",
        "        that have fewer actions (since -1 indicates \"pass\") and 'values' is set\n",
        "        to -inf on all pass moves.\n",
        "        \"\"\"\n",
        "        # actions_list is a list of lists of tuples. One list per game.\n",
        "        actions_list = game.get_available_actions()\n",
        "        num_actions_each_game = torch.as_tensor(\n",
        "                                        [len(acts) for acts in actions_list])\n",
        "        max_actions = num_actions_each_game.max()\n",
        "        # Allocate values output\n",
        "        values = float('-inf')*torch.ones(game.n_games, max_actions)\n",
        "        # Create a tensor of [n_games, max_actions, 2] that contains\n",
        "        # all available actions for all games, padded with -1s (pass move)\n",
        "        # where games have fewer than max_actions possible actions.\n",
        "        actions_tensor = -torch.ones(game.n_games, max_actions, 2,\n",
        "                                     dtype=torch.int32)\n",
        "        for g in range(game.n_games):\n",
        "            if num_actions_each_game[g] == 0:\n",
        "                continue\n",
        "            actions_tensor[g, :num_actions_each_game[g], :] =\\\n",
        "                torch.as_tensor(actions_list[g], dtype=torch.int32)\n",
        "        \n",
        "        # Evaluate each action, simultaneously for all games\n",
        "        for i in range(max_actions):\n",
        "            state = game.copy_state()\n",
        "            game.step(actions_tensor[:,i,:])\n",
        "            values[:,i] = self.evaluate(game, whoami)\n",
        "            values[num_actions_each_game < i+1] = float('-inf')\n",
        "            game.paste_state(state)\n",
        "        return values, actions_tensor, actions_list\n",
        "\n",
        "\n",
        "game = OthelloGame()\n",
        "agent = EpsilonGreedyAgent(epsilon=0.05)\n",
        "act = agent.select_move(game)\n",
        "game.step(act)\n",
        "game.render()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qsXnRsc-VP7B"
      },
      "source": [
        "Now **you** play a game against your $\\epsilon$-greedy agent"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UZM1_hoZVO-9"
      },
      "source": [
        "agent1 = \"human\"\n",
        "agent2 = EpsilonGreedyAgent(0.1)\n",
        "interface = InteractiveOthelloGame(n=6, player1=agent1, player2=agent2)\n",
        "interface.next_turn()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "emhImf3HxqGj"
      },
      "source": [
        "Next we'll create a function to pit two AI agents against each other"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eaYY4rcvLIun",
        "cellView": "form"
      },
      "source": [
        "#@markdown Define `ai_vs_ai`\n",
        "def ai_vs_ai(agent1, agent2=None, n_games=1, randomize_players=True,\n",
        "             print_result=False, progbar=True, n=8):\n",
        "    \"\"\"\n",
        "    Pit 2 AI players against each other (or self-play if just 1 given). Return\n",
        "    end-state values and list of board positions.\n",
        "    \"\"\"\n",
        "    if agent2 is None:\n",
        "        agent2 = agent1\n",
        "    game = OthelloGame(n_games=n_games, n=n)\n",
        "    saved_boards = [game.boards.clone()]\n",
        "    agents = [agent1, agent2]\n",
        "    if progbar:\n",
        "        # Games rarely go beyond n^2-4 moves since passes are rare\n",
        "        progress_bar = tqdm(total=game.n**2-4, desc='# moves per game')\n",
        "    while torch.any(game.game_status == OthelloGame.ACTIVE):\n",
        "        act = agents[game.current_player-1].select_move(game)\n",
        "        game.step(act)\n",
        "        saved_boards.append(game.boards.clone())\n",
        "        # Flag all boards that are done by flooding them with nan\n",
        "        saved_boards[-1][game.game_status==OthelloGame.OVER, ...] = float('nan')\n",
        "        if progbar: progress_bar.update(1)\n",
        "    if progbar:\n",
        "        progress_bar.total = progress_bar.n\n",
        "        progress_bar.close()\n",
        "    values = game.score_games()\n",
        "    if print_result:\n",
        "        for g in n_games:\n",
        "            print(f\"Game {g+1} of {n_games}:\", end=\"\\t\")\n",
        "            if values[0] == 0:\n",
        "                print(\"Draw\")\n",
        "            else:\n",
        "                print(f\"{agents[torch.argmax(values)].name}\"\n",
        "                \"(Player {torch.argmax(values)+1}) is the winner!\")\n",
        "    # Return end-game value according to player 1 and 2, as well as a list of\n",
        "    # all board states throughout play\n",
        "    return values, saved_boards"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XFlMeoUyVMoS"
      },
      "source": [
        "agent1 = EpsilonGreedyAgent(epsilon=0.5)\n",
        "agent2 = EpsilonGreedyAgent(epsilon=0.5)\n",
        "win_statistics(agent1, agent2, games_per_side=100);"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Vu59S6kDe5aP"
      },
      "source": [
        "---\n",
        "## Section 4.2: Training an NN to estimate value via self-play\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xoHV1UGAM1kr",
        "cellView": "form"
      },
      "source": [
        "#@title Video: How To Implement a DL Agent to Play Othello\n",
        "\n",
        "video = YouTubeVideo(id=\"MaW1WQ_O31w\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QfANPanbYmvp"
      },
      "source": [
        "\n",
        "Here we define a little network. We'll learn much more about this soon!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eJDAxf_FilrK"
      },
      "source": [
        "# Set up a network for position evaluation. \n",
        "class ValueNet(nn.Module):\n",
        "    def __init__(self, dropout=0.3, num_channels=64, n=8):\n",
        "        super(ValueNet, self).__init__()\n",
        "\n",
        "        # game params\n",
        "        self.n=n\n",
        "        self.num_channels = num_channels\n",
        "\n",
        "        # two conv layers\n",
        "        self.conv1 = nn.Conv2d(1,  num_channels, 3, stride=1)\n",
        "        self.conv2 = nn.Conv2d(num_channels, num_channels, 3, stride=1)\n",
        "\n",
        "        # two linear layers\n",
        "        self.fc1 = nn.Linear(num_channels*(self.n-4)**2, num_channels*2)\n",
        "        self.fc2 = nn.Linear(num_channels*2, 1)\n",
        "\n",
        "    def forward(self, s):\n",
        "        s = s.view(-1, 1, self.n, self.n)               \n",
        "        s = F.relu(self.conv1(s))                     \n",
        "        s = F.relu(self.conv2(s))                         \n",
        "        s = s.view(-1, self.num_channels*(self.n-4)**2)\n",
        "\n",
        "        s = F.relu(self.fc1(s))\n",
        "        v = self.fc2(s)                                                                       \n",
        "\n",
        "        return torch.tanh(v)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "finrdboHidsk"
      },
      "source": [
        "---\n",
        "## Section 4.3: Train the network!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nUQvR40RbFC4"
      },
      "source": [
        "### Getting training data\n",
        "\n",
        "Above, we defined the self-play function `ai_vs_ai`. Let's inspect its output."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YFWicfKmf1RB"
      },
      "source": [
        "agent1 = EpsilonGreedyAgent(epsilon=0.1)\n",
        "raw_outcomes, raw_states = ai_vs_ai(agent)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wVFQkD_Tf1y6"
      },
      "source": [
        "print(raw_outcomes, raw_states[0])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UkuONIWO987K"
      },
      "source": [
        "This pair of information will be the input and target for training our neural network below."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8T4gAmTP7mps"
      },
      "source": [
        "## Aside: preparing examples and data augmentation\n",
        "\n",
        "Because of symmetries in the board, each example can be *augmented* into many examples. Once you've seen one board and outcome pair, you've actually seen many! This code augments the data for more efficient training.\n",
        "\n",
        "If you want to learn more (and are way ahead of schedule), please check out the **Appendix** after the main tutorial.\n",
        "\n",
        "\n",
        "---\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gEbKl_zFNxRH",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Training A Network by Optimizing Loss\n",
        "\n",
        "video = YouTubeVideo(id=\"URVG1zyhaY0\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "asdTH5Wdq8cZ"
      },
      "source": [
        "The first thing we'll do is train a value network in a supervised learning setting. We'll observe games played out by some agents (your $\\epsilon$-greedy agents), and try to predict the outcome. This prediction is the value estimate.\n",
        "\n",
        "Below, we've defined a super standard train loop. We'll see loops like this over and over again this semester. It's important to walk through this code to get a sense of what happens."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dO3aZBeziUGX",
        "cellView": "form"
      },
      "source": [
        "#@markdown (RUN ME) Train helper functions `prepare_examples` and `boards2nn`\n",
        "SYMMETRIES = [lambda x: x,\n",
        "              lambda x: torch.rot90(x, 1, dims=(1, 2)),\n",
        "              lambda x: torch.rot90(x, 2, dims=(1, 2)),\n",
        "              lambda x: torch.rot90(x, 3, dims=(1, 2)),\n",
        "              lambda x: torch.flip(x, dims=(1,)),\n",
        "              lambda x: torch.flip(x, dims=(2,)),\n",
        "              lambda x: torch.transpose(x, 1, 2),\n",
        "              lambda x: torch.rot90(torch.flip(x,dims=(1,)), 1, dims=(1,2))]\n",
        "\n",
        "def boards2nn(board_states) -> torch.Tensor:\n",
        "    \"\"\"\n",
        "    Canonical preprocessing, taking in a set of boards and outputting NN inputs.\n",
        "    \"\"\"\n",
        "    if isinstance(board_states, list):\n",
        "        board_states = torch.cat(board_states, dim=0)\n",
        "    else:\n",
        "        board_states = board_states.clone() # Copy so not to affect the original\n",
        "    # Convert from [0,1,2] to [0,1,-1] (player 2 represented with -1)\n",
        "    board_states[torch.where(board_states == 2)] = -1\n",
        "    return board_states\n",
        "\n",
        "def prepare_examples(player1_value:torch.Tensor,\n",
        "                     board_states:List[torch.Tensor]):\n",
        "    \"\"\"Take raw game observations (final value from P1 perspective, list of\n",
        "    boards) and prepare a batch of training examples for the NN. Does 3 things\n",
        "    to the input:\n",
        "    1. Concatenates observations together\n",
        "    2. Adds \"inverse\" of all the inputs to account for player 2's perspective\n",
        "    3. Adds symmetries\n",
        "\n",
        "    In total, this makes the effective # examples 16 times as big as the number \n",
        "    of input positions (8 symmetries from each of 2 players' perspectives)\n",
        "    \n",
        "    player1_value is size [n_games]\n",
        "    board_states is size [max_num_moves].\n",
        "    Each entry contains a [n_games,n,n] board state for each game at that move.\n",
        "    \"\"\"\n",
        "    n_moves = len(board_states)\n",
        "    n_games, n, _ = board_states[0].size()\n",
        "    # Concatenate and preprocess board states.\n",
        "    # Result is size [n_games*n_moves,n,n]\n",
        "    # ordered like [game0turn0, game1turn0, ... gameNturn0,\n",
        "    #               game0turn1, game1turn1, ... ]\n",
        "    board_states = boards2nn(board_states)\n",
        "    # Repeat all the player-1 values once per move. The input value are already\n",
        "    # ordered like [game0, game1, ..., gameN] so to match the board order we\n",
        "    # simply repeat() the values\n",
        "    outcomes = player1_value.repeat(n_moves).float()\n",
        "    # It may be the case that not all games ran for n_moves. These are indicated\n",
        "    # by nan values in the board state (see ai_vs_ai)\n",
        "    drop_states = torch.isnan(board_states[:,0,0])\n",
        "    board_states = board_states[~drop_states]\n",
        "    outcomes = outcomes[~drop_states]\n",
        "    # Add all symmetries (increases effective data 8-fold)\n",
        "    outcomes = outcomes.repeat(len(SYMMETRIES))\n",
        "    board_states = torch.cat([sym(board_states) for sym in SYMMETRIES], dim=0)\n",
        "    # Add in examples from player 2's perspective by flipping signs on both\n",
        "    # value and on the board state itself\n",
        "    outcomes = torch.cat([outcomes, -outcomes], dim=0)\n",
        "    board_states = torch.cat([board_states, -board_states], dim=0)\n",
        "    return outcomes, board_states"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ov_xne5Xx-04"
      },
      "source": [
        "def train_value_by_self_play(value_net, agent, n_iters=50, games_per_iter=32,\n",
        "                             n=8):\n",
        "    \"\"\"\n",
        "    Play multiple batches of games where 'agent' plays against itself. The board\n",
        "    positions and game outcomes will be used to train the neural net 'value_net'\n",
        "    \"\"\"\n",
        "    # first we define the optimizer and initialize it with a learning rate. \n",
        "    # We're using Adam, but we could also have used stochastic gradient descent\n",
        "    # via torch.optim.SGD\n",
        "    opt = torch.optim.Adam(value_net.parameters(), lr=1e-3)\n",
        "    # we'll minimize the MSE from the outcome to the network's output\n",
        "    v_criterion = nn.MSELoss()\n",
        "\n",
        "    # we'll loop over a number of network updates, each of which includes\n",
        "    # a batch of self-play games\n",
        "    losses = torch.zeros(n_iters).cuda()\n",
        "    for i in tqdm(range(n_iters), desc='epochs'):\n",
        "        ##### Here we build our training examples #####\n",
        "        # batch some games and train based on player 1's perspective\n",
        "        outcomes, board_states = [], []\n",
        "        # Play AI against itself, running games_per_iter games all in parallel\n",
        "        raw_outcomes, raw_states = ai_vs_ai(agent, n=n, n_games=games_per_iter,\n",
        "                                            progbar=False)\n",
        "        # Get all trainable information from this batch of games\n",
        "        outcomes, board_states = prepare_examples(raw_outcomes[:,0], raw_states)\n",
        "\n",
        "        ##### Here we train the network\n",
        "        value_net.train() # Ensure we are in 'training' mode, not 'evaluation'\n",
        "        opt.zero_grad() # ready the optimizer\n",
        "        estimated_v = value_net(board_states).view(-1)  # predict\n",
        "        loss = v_criterion(outcomes, estimated_v) # loss\n",
        "        loss.backward() # get gradients of loss w/r/t the network parameters\n",
        "        opt.step() # take a step to lower the loss\n",
        "\n",
        "        losses[i] = loss.detach()\n",
        "\n",
        "    return losses.cpu().numpy()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PodjvTJZOfJk"
      },
      "source": [
        "Create a new value network and train it on the output of the $\\epsilon$-greedy player's games"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8cHhVgRl1LJe"
      },
      "source": [
        "value_net = ValueNet(n=6).cuda()\n",
        "heuristic_agent = EpsilonGreedyAgent(0.25, name=\"heuristic\")\n",
        "mse = train_value_by_self_play(value_net, heuristic_agent, n=6)\n",
        "\n",
        "plt.figure()\n",
        "plt.plot(mse)\n",
        "plt.title(f\"Training on self-play games by {heuristic_agent.name}\")\n",
        "plt.xlabel('Iteration #')\n",
        "plt.ylabel('MSE loss')\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5aClEmIteiJb"
      },
      "source": [
        "**Student response:** Here we're minimizing the MSE between the outcome and the NN's estimate. Is that the best choice of loss criterion? Discuss with your group if this makes sense or if there's a better alternative."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FsEOGTPie9W3",
        "cellView": "form"
      },
      "source": [
        "better_criterion_or_no = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R5MGWI2XdmrF"
      },
      "source": [
        "Let's make sure the value estimate makes sense. We've made a plotting routine that, for one batch of games, outputs the value estimate during the games. The lines are colored by the eventual true outcome."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fr3QbUC3V--C"
      },
      "source": [
        "def plot_value_predictions(value_net, agent, n_games=100, title=None):\n",
        "    value_net.eval() # Ensure 'evaluation' rather than 'training' mode\n",
        "\n",
        "    plt.figure(figsize=(8,4))\n",
        "    ax1 = plt.subplot(1,2,1)\n",
        "    ax2 = plt.subplot(1,2,2)\n",
        "    outcomes, boards = ai_vs_ai(agent, n_games=n_games, progbar=True, n=6)\n",
        "    board_nn = boards2nn(boards)\n",
        "    # Note: some of boards_nn will be nan if the games ended early, in which\n",
        "    # case we don't want to plot them anyway. Force those outputs to nan too.\n",
        "    predicted_values_1 = value_net(board_nn)\n",
        "    predicted_values_2 = value_net(-board_nn)\n",
        "    predicted_values_1[torch.isnan(board_nn[:,0,0])] = float('nan')\n",
        "    predicted_values_2[torch.isnan(board_nn[:,0,0])] = float('nan')\n",
        "    for i in range(n_games):\n",
        "        vals_1_this_game = predicted_values_1[i::n_games]\n",
        "        vals_2_this_game = predicted_values_2[i::n_games]\n",
        "        ax1.plot(vals_1_this_game.detach().cpu(),\n",
        "                 color=outcome2color(outcomes[i,0].item(), 0.1))\n",
        "        ax2.plot(vals_2_this_game.detach().cpu(),\n",
        "                 color=outcome2color(outcomes[i,1].item(), 0.1))\n",
        "    for i, ax in enumerate([ax1, ax2]):\n",
        "        ax.set_xlabel('Turn #')\n",
        "        ax.set_ylabel('Predicted Outcome')\n",
        "        ax.grid()\n",
        "        ax.set_title(f\"Player {i+1}'s perspective\")\n",
        "    if title is not None:\n",
        "        plt.suptitle(title, y=1.2, fontsize=14)\n",
        "    plt.tight_layout()\n",
        "    plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oiI_EJ_YdYEo"
      },
      "source": [
        "plot_value_predictions(value_net, heuristic_agent, 100, title=f\"NN predictions for {heuristic_agent.name} self-play\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3_HGiTAqf7FW"
      },
      "source": [
        "Discuss with your group as to whether this makes sense."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sDwGk4_3e4DB"
      },
      "source": [
        "---\n",
        "# Section 5: Making a neural $\\epsilon$-greedy player\n",
        "\n",
        "Now that we have a neural net that does a decent job of predicting game outcomes, we can use the neural network itself to evaluate the board and make greedy decisions! Let's try it.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5H8O_ZeOPWYi",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Making Sense of the Value Function\n",
        "\n",
        "param_list = ['better_criterion_or_no']\n",
        "for param in param_list:\n",
        "    if param not in locals():\n",
        "        raise NameError(\"Please make sure to run the cell \"\n",
        "                        \"after submitting your answer for <{}>!\".format(param))\n",
        "\n",
        "video = YouTubeVideo(id=\"7s0aMrNFExo\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "try: t5;\n",
        "except NameError: t5=time.time()\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SbaUyPR6cRGf"
      },
      "source": [
        "# Create a new epsilon-greedy agent whose value function comes from neural net\n",
        "class NeuralGreedyAgent(EpsilonGreedyAgent):\n",
        "    def __init__(self, value_net, epsilon=0.05, name=None):\n",
        "        self.value_net = value_net\n",
        "        name = name if name is not None else f\"NeuralEpsilon{epsilon}\"\n",
        "        super(NeuralGreedyAgent, self).__init__(name=name, epsilon=epsilon)\n",
        "\n",
        "    def evaluate(self, game:OthelloGame, whoami:int) -> float:\n",
        "        \"\"\"NeuralGreedyAgent.evaluate overrides EpsilonGreedyAgent.evaluate so\n",
        "        that EpsilonGreedyAgent.select_move logic is unchanged.\n",
        "        \"\"\"\n",
        "        # Ensure network is in 'evaluation' rather than 'training' mode\n",
        "        self.value_net.eval()\n",
        "        # Run the neural net on the board (player 1's perspective by default)\n",
        "        v = self.value_net(boards2nn(game.boards)).flatten()\n",
        "        # Grab the output as a float, but flip sign if \"I\" am player 2, since\n",
        "        # the value net was always trained from the perspective of the player\n",
        "        # whose pieces are \"+1\" on the board.\n",
        "        return v if whoami == 1 else -v"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8cmMsY_KfitP"
      },
      "source": [
        "Did we improve? Let's see who wins between the 'Neural' agent and the 'Heuristic' agent we hand-coded earlier. Remember: the neural net is simply trained to predict who wins __when the heuristic agent plays against itself!__"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UVNueX-4fhqB"
      },
      "source": [
        "neural_agent = NeuralGreedyAgent(value_net, epsilon=0.25)\n",
        "win_statistics(neural_agent, heuristic_agent)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JE0y3qIMhdqp"
      },
      "source": [
        "---\n",
        "# Section 6: Lifting our value estimate up by its bootsraps"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TOO9SLjhCHjs",
        "cellView": "form"
      },
      "source": [
        "#@markdown I'm a time tracker, please run me.\n",
        "try: t6;\n",
        "except NameError: t6=time.time()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mn9pwPiMfzVi"
      },
      "source": [
        "Perhaps we can do even better? Rather than predicting the outcome of the Heuristic agent's games, why don't we train the neural net to __predict the outcomes of its own games.__\n",
        "\n",
        "To do this, the network will be alternating between two modes: in `evaluation` mode (`value_net.eval()`) it is just an input-output function, giving us its estimate of the value of board positions. After every batch of games, we switch into `training` mode (`value_net.train()`), replay all the games that just happened, and try to minimize prediction errors on them."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nqCASsHU-lR5"
      },
      "source": [
        "# Further train the value net through self-play\n",
        "# (but make a copy of the net so we can compare before and after)\n",
        "value_net_2 = deepcopy(value_net).cuda()\n",
        "self_play_agent = NeuralGreedyAgent(value_net_2, epsilon=0.25)\n",
        "# Train value_net_2 using games where the self_play_agent plays against itself\n",
        "self_play_mse = train_value_by_self_play(value_net_2, self_play_agent, n=6)\n",
        "plt.figure()\n",
        "plt.title(\"Mean Square Error for self-play\")\n",
        "plt.plot(self_play_mse)\n",
        "plt.xlabel('Iteration #')\n",
        "plt.ylabel('MSE loss')\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ElCyEjxuhv5j"
      },
      "source": [
        "Why doesn't this appear to go down? Discuss with your group."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cvKZKe5lcWdw"
      },
      "source": [
        "plot_value_predictions(value_net_2, self_play_agent, 100,\n",
        "                       title=\"NN predictions for self-play\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vpDz5HWqmkjZ"
      },
      "source": [
        "We now have 3 agents: the original hand-coded \"heuristic\" agent, the \"neural\" agent that was trained to predict outcomes of the heuristic agent's games, and a \"self-play\" agent that further modified the \"neural\" agent's value net with games against itself.\n",
        "\n",
        "Who wins?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VAOeeh-1mbky"
      },
      "source": [
        "neural_agent.name = 'Super'         # Short for \"supervised training\"\n",
        "self_play_agent.name = 'Self-Super' # \"self-play after supervised pretraining\"\n",
        "heuristic_agent.name = 'Heuristic'  # The hand-coded heuristic agent\n",
        "\n",
        "win_statistics(self_play_agent, neural_agent)\n",
        "win_statistics(self_play_agent, heuristic_agent)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-08cIGGzlxoi"
      },
      "source": [
        "### (If time) play a game against your best agent"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mODlAguUlyGQ"
      },
      "source": [
        "interface = InteractiveOthelloGame(player1='human',\n",
        "                                   player2=self_play_agent, n=6)\n",
        "interface.next_turn()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0JfRgx4zcaU1"
      },
      "source": [
        "# Submit your responses\n",
        "Please run the following cell and then press \"Submit\" so we can record your responses."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0Ky8GBy6SBe1",
        "cellView": "form"
      },
      "source": [
        "#@title Video: The End of Part 1\n",
        "\n",
        "video = YouTubeVideo(id=\"m4k4Mo-wBj4\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "try:t7;\n",
        "except NameError: t7 = time.time()\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dGy0Zz6ccab2",
        "cellView": "form"
      },
      "source": [
        "import time\n",
        "import numpy as np\n",
        "from IPython.display import IFrame\n",
        "#@markdown #Run Cell to Show Airtable Form\n",
        "#@markdown ##**Confirm your answers and then click \"Submit\"**\n",
        "\n",
        "def prefill_form(src, fields: dict):\n",
        "  '''\n",
        "  src: the original src url to embed the form\n",
        "  fields: a dictionary of field:value pairs,\n",
        "  e.g. {\"pennkey\": my_pennkey, \"location\": my_location}\n",
        "  '''\n",
        "  prefills = \"&\".join([\"prefill_%s=%s\"%(key, fields[key]) for key in fields])\n",
        "  src = src + prefills\n",
        "  src = \"+\".join(src.split(\" \"))\n",
        "  return src\n",
        "\n",
        "\n",
        "#autofill time if it is not present\n",
        "try: t0;\n",
        "except NameError: t0 = time.time()\n",
        "try: t1;\n",
        "except NameError: t1 = time.time()\n",
        "try: t2;\n",
        "except NameError: t2 = time.time()\n",
        "try: t3;\n",
        "except NameError: t3 = time.time()\n",
        "try: t4;\n",
        "except NameError: t4 = time.time()\n",
        "try: t5;\n",
        "except NameError: t5 = time.time()\n",
        "try: t6;\n",
        "except NameError: t6 = time.time()\n",
        "try: t7;\n",
        "except NameError: t7 = time.time()\n",
        "\n",
        "#autofill fields if they are not present\n",
        "#a missing pennkey and pod will result in an Airtable warning\n",
        "#which is easily fixed user-side.\n",
        "try: my_pennkey;\n",
        "except NameError: my_pennkey = \"\"\n",
        "try: my_pod;\n",
        "except NameError: my_pod = \"Select\"\n",
        "try: what_i_want;\n",
        "except NameError: what_i_want = \"\"\n",
        "try: theory_or_implementation;\n",
        "except NameError: theory_or_implementation = 5\n",
        "try: play_analysis;\n",
        "except NameError: play_analysis = \"\"\n",
        "try: network_role;\n",
        "except NameError: network_role = \"\"\n",
        "try: how_to_train_value;\n",
        "except NameError: how_to_train_value = \"\"\n",
        "try: week\n",
        "except NameError: week = \"\"\n",
        "try: epsilon_effect\n",
        "except NameError: epsilon_effect = \"\"\n",
        "try: better_criterion_or_no\n",
        "except NameError: better_criterion_or_no = \"\"\n",
        "\n",
        "\n",
        "times = np.array([t1,t2,t3,t4,t5,t6,t7])-t0\n",
        "\n",
        "fields = {\"pennkey\": my_pennkey,\n",
        "          \"pod\": my_pod,\n",
        "          \"expectation\":what_i_want,\n",
        "          \"theory_imp\": theory_or_implementation,\n",
        "          \"human_play_analysis\":play_analysis,\n",
        "          \"network_role\":network_role,\n",
        "          \"how_to_train_value\":how_to_train_value,\n",
        "          \"week_most_excited\":week,\n",
        "          \"better_criterion_or_no\":better_criterion_or_no,\n",
        "          \"cumulative_times\": times}\n",
        "\n",
        "src = \"https://airtable.com/embed/shrUTMABKzfiDFEON?\"\n",
        "\n",
        "#now instead of the original source url, we do: src = prefill_form(src, fields)\n",
        "display(IFrame(src = prefill_form(src, fields), width = 800, height = 400))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "00BtYL678gKb"
      },
      "source": [
        "# Appendix\n",
        "\n",
        "\n",
        "If you have extra time, come here to learn a little about symmetry in data and how to use that to train more efficiently. This material will also be revisted later in the course."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6eFryptM8r5N"
      },
      "source": [
        "### Exercise\n",
        "Briefly, before you continue, by what factor can we increase our training set size *without running any games?* "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nqjGcAF18r5W",
        "cellView": "form"
      },
      "source": [
        "free_data_augmentation_factor = 0 #@param {type:\"number\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PcZg1SDT8r5W"
      },
      "source": [
        "Below, we've reproduced some functions from above that perform data augmentation from the results of `ai_vs_ai`. They also convert the data from numpy into pytorch."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bikGs6Nj8r5X",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Data Augmentation and Batching\n",
        "\n",
        "video = YouTubeVideo(id=\"F9wx_GdXe4Q\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Wczzrutb8r5W"
      },
      "source": [
        "# can you make sense of this list of symmetries?\n",
        "SYMMETRIES = [lambda x: x,\n",
        "              lambda x: torch.rot90(x, 1, dims=(1, 2)),\n",
        "              lambda x: torch.rot90(x, 2, dims=(1, 2)),\n",
        "              lambda x: torch.rot90(x, 3, dims=(1, 2)),\n",
        "              lambda x: torch.flip(x, dims=(1,)),\n",
        "              lambda x: torch.flip(x, dims=(2,)),\n",
        "              lambda x: torch.transpose(x, 1, 2),\n",
        "              lambda x: torch.rot90(torch.flip(x,dims=(1,)), 1, dims=(1,2))]\n",
        "\n",
        "def boards2nn(board_states) -> torch.Tensor:\n",
        "    \"\"\"\n",
        "    Canonical preprocessing, taking in a set of boards and outputting NN inputs.\n",
        "    \"\"\"\n",
        "    if isinstance(board_states, list):\n",
        "        board_states = torch.cat(board_states, dim=0)\n",
        "    else:\n",
        "        board_states = board_states.clone() # Copy so not to affect the original\n",
        "    # Convert from [0,1,2] to [0,1,-1] (player 2 represented with -1)\n",
        "    board_states[torch.where(board_states == 2)] = -1\n",
        "    return board_states\n",
        "\n",
        "def prepare_examples(player1_value:torch.Tensor,\n",
        "                     board_states:List[torch.Tensor]):\n",
        "    \"\"\"Take raw game observations (final value from P1 perspective, list of\n",
        "    boards) and prepare a batch of training examples for the NN. Does 3 things\n",
        "    to the input:\n",
        "    1. Concatenates observations together\n",
        "    2. Adds \"inverse\" of all the inputs to account for player 2's perspective\n",
        "    3. Adds symmetries\n",
        "\n",
        "    In total, this makes the effective # examples 16 times as big as the number \n",
        "    of input positions (8 symmetries from each of 2 players' perspectives)\n",
        "    \n",
        "    player1_value is size [n_games]\n",
        "    board_states is size [max_num_moves].\n",
        "    Each entry contains a [n_games,n,n] board state for each game at that move.\n",
        "    \"\"\"\n",
        "    n_moves = len(board_states)\n",
        "    n_games, n, _ = board_states[0].size()\n",
        "    # Concatenate and preprocess board states.\n",
        "    # Result is size [n_games*n_moves,n,n]\n",
        "    # ordered like [game0turn0, game1turn0, ... gameNturn0,\n",
        "    #               game0turn1, game1turn1, ... ]\n",
        "    board_states = boards2nn(board_states)\n",
        "    # Repeat all the player-1 values once per move. The input value are already\n",
        "    # ordered like [game0, game1, ..., gameN] so to match the board order we\n",
        "    # simply repeat() the values\n",
        "    outcomes = player1_value.repeat(n_moves).float()\n",
        "    # It may be the case that not all games ran for n_moves. These are indicated\n",
        "    # by nan values in the board state (see ai_vs_ai)\n",
        "    drop_states = torch.isnan(board_states[:,0,0])\n",
        "    board_states = board_states[~drop_states]\n",
        "    outcomes = outcomes[~drop_states]\n",
        "    # Add all symmetries (increases effective data 8-fold)\n",
        "    outcomes = outcomes.repeat(len(SYMMETRIES))\n",
        "    board_states = torch.cat([sym(board_states) for sym in SYMMETRIES], dim=0)\n",
        "    # Add in examples from player 2's perspective by flipping signs on both\n",
        "    # value and on the board state itself\n",
        "    outcomes = torch.cat([outcomes, -outcomes], dim=0)\n",
        "    board_states = torch.cat([board_states, -board_states], dim=0)\n",
        "    return outcomes, board_states"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pu9Fi5uw86Vn"
      },
      "source": [
        "## Inspecting the `prepare_examples` code."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MwZde3gt8r5Y"
      },
      "source": [
        "heuristic_agent = EpsilonGreedyAgent(0.25, name=\"heuristic\")\n",
        "\n",
        "outcomes, board_states = [], []\n",
        "# Play AI against itself, running games_per_iter games all in parallel\n",
        "raw_outcomes, raw_states = ai_vs_ai(agent, n_games=50)\n",
        "# Get all trainable information from this batch of games\n",
        "outcomes, board_states = prepare_examples(raw_outcomes[:,0], raw_states)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XwXZtw-e8r5Y"
      },
      "source": [
        "Inspect the size of `board_states` and `outcomes`:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9j3AZ99S8r5Y"
      },
      "source": [
        "print(outcomes.size(), board_states.size())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y6MFxOzH8r5Y"
      },
      "source": [
        "That's a big number! Where did that come from? \n",
        "\n",
        "**Student response:** Describe precisely why these objects have those sizes. What factors are multiplied to produce this number of boards, and what is their value?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wOPVNxh-8r5Y",
        "cellView": "form"
      },
      "source": [
        "why_this_size = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}